"""
Bounding Volume Hierarchies
"""
# pylint: disable=wildcard-import,unused-wildcard-import,trailing-whitespace
# pylint: disable=consider-using-f-string,no-value-for-parameter
import numpy as np
import taichi as ti
from .type_enum import (
    PrimitiveType,
    MaterialType,
    LightType
)
from .gmath import normalized,identity_mat4
from .gtaichi import *
from .light import *

TEX_WIDTH = 64*64

@ti.data_oriented
class AabbBounding:
    """
    扁平的AAbb包围盒结构
    """
    def __init__(self,scene):
        self.scene = scene
        self.obj2material = {}  #scene对象id映射到material
        self.name2tex = {}      #材质名称到材质索引映射    
        self.m2n = {}  
        self.sky_i = -1
        self.alloc()
        self.load_material()
        self.load_primitive()

    def get_scene_info(self):
        """
        取得场景中的对象数量,三角形数量,材质数量,顶点数量
        """
        obj_n,tran_n,tex_n,sphere_n,light_n = 0,0,0,0,0
        obj_n = len(self.scene.objs)
        for obj in self.scene.objs:
            if "texture" and obj.material:
                tex_n += 1
            if obj.mesh and 'trangles' in obj.mesh:
                if obj.mesh['type']=='sphere':
                    sphere_n += 1
                else:
                    tran_n += obj.mesh['trangles'].shape[0]
            if obj.light:
                if obj.light['type']==LightType.SKY_LIGHT:
                    sphere_n += 1
                light_n += 1
        return obj_n,tran_n,tex_n,sphere_n,light_n
    
    def alloc(self):
        """
        重新分配全部数据
        """
        # pylint: disable=no-member
        obj_n,tran_n,tex_n,sphere_n,light_n = self.get_scene_info()

        self.bounds = AabbBound.field(shape=(obj_n+1,))
        self.sphere = Sphere.field(shape=(sphere_n+1,))
        self.materials = Material.field(shape=(obj_n+1,))
        self.light = Light.field(shape=(light_n+1,))
        self.trangle = Trangle.field(shape=(tran_n,))
        self.np_trangle = np.ndarray(shape=(tran_n,18),dtype=np.float32)
    
        self.textures = ti.Vector.field(3,dtype=ti.f32)
        ti.root.dense(ti.i,(tex_n+1,)).pointer(ti.jk,(64,64)).\
            dense(ti.jk,(64,64)).place(self.textures)
        #self.textures = ti.Vector.field(3,dtype=ti.f32,shape=(tex_n,TEX_WIDTH,TEX_WIDTH))
        self.texture_size = ti.Vector.field(2,dtype=ti.i32,shape=(tex_n+1,))
        self.tex_n = 0
        self.tran_n = 0
        self.sphere_n = 0
        self.bound_n = 0
        self.light_n = 0
        self.material_n = 0
        self.sky_i = -1
        self.environment_color = ti.Vector.field(3,dtype=ti.f32,shape=())
        self.environment_color[None] = ti.Vector([0.6,0.6,0.6])        

    def load_material(self):
        """
        重新加载材质
        """
        i = 0
        for obj in self.scene.objs:
            if obj.light:
                continue
            m = obj.material
            if m["texture"]:
                if m["texture"] in self.name2tex:
                    tex = self.name2tex[m["texture"]]
                else:
                    tex = self.load_tex(m["texture"])
            else:
                tex = -1
            t = m["type"]
            color = m["color"]
            a = max(m["a"],0.02)
            if obj.light is None:
                self.materials[i] = Material(t.value,
                                             ti.Vector(color),
                                             tex,
                                            a,
                                            ti.Vector(m["f"]),
                                            m["metallic"],
                                            m["mu"])
                self.obj2material[id(obj)] = i
                self.m2n[i] = [obj.name,obj.mesh["type"],0,obj]
                i += 1
        self.material_n = i

    @ti.kernel
    def copy_tex(self,tex:int,img:ti.types.ndarray()):# pylint: disable=no-value-for-parameter
        """
        将numpy数组图像复制到对应的纹理,剪切掉大于TEX_WIDTH的部分
        """
        for i in range(min(TEX_WIDTH,img.shape[0])):
            for j in range(min(TEX_WIDTH,img.shape[1])):
                self.textures[tex,i,j] = ti.Vector([img[i,j,0]/255.,
                                                    img[i,j,1]/255.,
                                                    img[i,j,2]/255.])

    def load_tex(self,name):
        """
        加载纹理图name,返回纹理索引
        """
        if name in self.name2tex:
            return self.name2tex[name]
        if name:
            tex = self.tex_n
            img = ti.tools.imread(name)
            if img is not None:
                self.texture_size[tex][0] = img.shape[0]
                self.texture_size[tex][1] = img.shape[1]
                img_2 = np.ndarray(dtype=np.uint8,shape=img.shape)
                img_2[:] = img[:]
                self.copy_tex(tex,img_2)
                self.name2tex[name] = tex
                self.tex_n+=1
                return tex
        return -1

    def add_sphere(self,mat,obj,add_bound=True):
        """
        加入一个球对象
        """
        if obj.light and obj.light['type']==LightType.SKY_LIGHT:
            r = SKY_SPHERE_RADIUS
        else:
            r = obj.mesh['meta'][0]
        center = mat @ np.array([0.,0.,0.,1.])
        x_axis = mat @ np.array([1.,0.,0.,1.])
        y_axis = mat @ np.array([0.,1.,0.,1.])
        z_axis = mat @ np.array([0.,0.,1.,1.])
        self.sphere[self.sphere_n] = Sphere( ti.Vector(center[:3]),
                                        r,
                                        ti.Vector(x_axis[:3]),
                                        ti.Vector(y_axis[:3]),
                                        ti.Vector(z_axis[:3]))
        cp = ti.Vector(center[:3])
        vr = ti.Vector((r,r,r))
        if add_bound:
            m = self.obj2material[id(obj)]
            self.m2n[m][2] = self.bound_n
            self.bounds[self.bound_n] = AabbBound(PrimitiveType.SPHERE.value,
                                                    cp-vr,
                                                    cp+vr,
                                                    self.sphere_n,
                                                    0,
                                                    m)
            self.bound_n+=1
        self.sphere_n+=1
        return self.sphere_n-1
    
    def add_mesh(self,mat,obj,np_tran,add_bound=True):
        """
        加入一个网格
        """
        mesh = obj.mesh
        points = (mat @ mesh['points'].T).T
        trangle = mesh['trangles']
        tex_coord = mesh['textures']
        norm_n = trangle.shape[0]
        if "normales" in mesh:
            norm_n = mesh['normales'].shape[0]
        norm = np.empty(shape=(norm_n,3),dtype=np.float32)
        if "normales" in mesh:
            rmat4 = np.copy(mat)
            rmat4[:3,3] = 0
            rmat4[3,:3] = 0
            normales = (rmat4 @ mesh['normales'].T).T
            norm[:,:] = normalized(normales[:,:3])
        else:
            for i in range(trangle.shape[0]):
                a = points[trangle[i][0],:3]
                b = points[trangle[i][1],:3]
                c = points[trangle[i][2],:3]
                norm[i,:] = normalized(np.cross(c-a,b-a))
        amin = points[0][:3]
        amax = points[0][:3]
        for i in range(trangle.shape[0]):
            a = points[trangle[i][0],:3]
            b = points[trangle[i][1],:3]
            c = points[trangle[i][2],:3]
            amin = np.fmin(amin,np.fmin(np.fmin(a,b),c))
            amax = np.fmax(amax,np.fmax(np.fmax(a,b),c))
            if trangle.shape[1]==6:
                a_uv = tex_coord[trangle[i][3],:2]
                b_uv = tex_coord[trangle[i][4],:2]
                c_uv = tex_coord[trangle[i][5],:2]
            else:
                a_uv = tex_coord[trangle[i][0],:2]
                b_uv = tex_coord[trangle[i][1],:2]
                c_uv = tex_coord[trangle[i][2],:2]
            if trangle.shape[1]==9:
                n = norm[trangle[i][6]]
            else:
                n = norm[i]
            np_tran[self.tran_n+i,:] = (*a,*b,*c,*a_uv,*b_uv,*c_uv,*n)
        if add_bound:
            m = self.obj2material[id(obj)]
            self.m2n[m][2] = self.bound_n
            self.bounds[self.bound_n] = AabbBound( PrimitiveType.TRANGLE.value,
                                                    amin,
                                                    amax,
                                                    self.tran_n,
                                                    trangle.shape[0],
                                                    m)
            self.bound_n += 1
        self.tran_n += trangle.shape[0]
        return self.tran_n-trangle.shape[0],trangle.shape[0]

    def add_light(self,mat,obj,np_tran):
        """
        加入灯
        """
        light = obj.light
        center = mat @ np.array([0.,0.,0.,1.])
        rmat4 = np.copy(mat)
        rmat4[:3,3] = 0
        rmat4[3,:3] = 0
        dir_ = rmat4 @ np.array([*light['dir'],1.])
        offset,count = -1,0
        if light["texture"]:
            if light["texture"] in self.name2tex:
                tex = self.name2tex[light["texture"]]
            else:
                tex = self.load_tex(light["texture"])
        else:
            tex = -1
        if light['type']==LightType.AREA_LIGHT or light['type']==LightType.SKY_LIGHT:
            if light['type']==LightType.SKY_LIGHT:
                offset = self.add_sphere(mat,obj,False)
                self.sky_i = offset
            elif light['type']==LightType.AREA_LIGHT and obj.mesh['type'] == 'sphere':
                offset = self.add_sphere(mat,obj,False)
            else:
                offset,count = self.add_mesh(mat,obj,np_tran,False)
        self.light[self.light_n] = Light(light['type'].value,
                                         light['power'],
                                         tex,
                                         ti.Vector(light['color']),
                                         ti.Vector(center[:3]),
                                         ti.Vector(normalized(dir_[:3])),
                                         light['angle'],
                                         offset,
                                         count,
                                         0.,
                                         0.)
        self.light_n += 1

    def load_primitive(self):
        """
        加载顶点,法线,三角形,贴图坐标
        """
        self.environment_color[None] = ti.Vector(self.scene.environment_color)
        np_tran = np.ndarray((self.trangle.shape[0],18),dtype=np.float32)
        def proccess_geometry(mat,o):
            mesh = o.mesh
            if o.light:
                self.add_light(mat,o,np_tran)
            elif mesh['type'] == 'sphere':
                self.add_sphere(mat,o)
            else:
                self.add_mesh(mat,o,np_tran)

        self.scene.for_each(proccess_geometry,identity_mat4())
        self.copy_trangle(np_tran)
        self.update_light()
    
    def update_light(self):
        """
        更新灯CDF
        """
        pow_sum = 0.
        v = 0.
        
        for i in range(self.light_n):
            pow_sum += self.light[i].power
        for i in range(self.light_n):
            light = self.light[i]
            light.pdf_v = light.power/pow_sum
            v += light.pdf_v
            light.cdf_v = v

    @ti.kernel
    def copy_trangle(self,x:ti.types.ndarray()):
        """
        批量将三角形数据传入显存
        """
        for i in range(self.trangle.shape[0]):
            self.trangle[i] = Trangle(
                                    ti.Vector([x[i,0],x[i,1],x[i,2]]), #a
                                    ti.Vector([x[i,3],x[i,4],x[i,5]]), #b
                                    ti.Vector([x[i,6],x[i,7],x[i,8]]), #c
                                    ti.Vector([x[i,9],x[i,10]]),
                                    ti.Vector([x[i,11],x[i,12]]),
                                    ti.Vector([x[i,13],x[i,14]]),
                                    ti.Vector([x[i,15],x[i,16],x[i,17]]))

    @ti.func   
    def intersection(self,ray:Ray):
        """
        射线和场景中的对象求交点
        """
        type_:ti.u32 = PrimitiveType.SPHERE.value
        b, t, m, l = 0, 1e38,0, 0
        p = ti.Vector([.0,.0,.0])
        n = ti.Vector([.0,.0,.1])
        for i in range(self.bound_n):
            bound = self.bounds[i]
            if ray_aabb_intersection(ray,bound.min,bound.max):
                if bound.type==PrimitiveType.TRANGLE.value:
                    for j in range(bound.offset,bound.offset+bound.count):
                        trangle = self.trangle[j]
                        b2,t2,p2,n2 = cross_ray_trangle(ray,trangle)
                        if b2 and t>t2:
                            type_ = bound.type
                            b,t,p,n,m,l = b2,t2,p2,n2,bound.material,j
                elif bound.type==PrimitiveType.SPHERE.value:
                    sphere = self.sphere[bound.offset]
                    b2,t2,p2,n2 = cross_ray_sphere(ray,sphere)
                    if b2 and t>t2:
                        type_ = bound.type
                        b,t,p,n,m,l = b2,t2,p2,n2,bound.material,bound.offset
                        
        return b,Intersection(type_,t,p,n,m,l)

    @ti.func
    def get_material(self,i:ti.i32)->Material:
        """
        返回索引为i的材质
        """
        return self.materials[i]
    
    @ti.func
    def get_uv_tex(self,tex,u,v):
        """
        取得材质tex的uv颜色
        """
        size = self.texture_size[tex]
        u = ti.max(0,ti.min(u,1.))
        v = ti.max(0,ti.min(v,1.))
        x = int((size[0]-1) * u)
        y = int((size[1]-1) * v)
        return self.textures[tex,x,y]
      
    @ti.func
    def get_trangle_tex(self,p,l,tex):
        """
        返回p点材质颜色,l为三角索引,tex为材质的索引
        """
        tran = self.trangle[l]
        ab = tran.b - tran.a
        ac = tran.c - tran.a
        ap = p - tran.a
        sx = 0.
        sy = 0.
        d = ab.x*ac.y-ab.y*ac.x
        if ti.abs(d)>EPS:
            sx = (ab.x*ap.y-ab.y*ap.x)/d
        else:
            d = ac.y*ab.z-ac.z*ab.y
            if ti.abs(d)>EPS:
                sx = (ap.y*ab.z-ap.z*ab.y)/d
            else:
                d = ac.x*ab.z-ac.z*ab.x
                if ti.abs(d)>EPS:
                    sx = (ap.x*ab.z-ap.z*ab.x)/d
        if ti.abs(ab.x)>EPS:
            sy = (ap.x-sx*ac.x)/ab.x
        elif ti.abs(ab.y)>EPS:
            sy = (ap.y-sx*ac.y)/ab.y
        else:
            sy = (ap.z-sx*ac.z)/ab.z
        p_uv = sx*(tran.c_uv-tran.a_uv) + sy*(tran.b_uv-tran.a_uv) + tran.a_uv
        return self.get_uv_tex(tex,p_uv.x,p_uv.y)

    @ti.func
    def get_sphere_tex(self,p:vec3,i,tex:ti.i32):
        """
        返回p点,i为球索引,tex为材质索引
        """
        sphere = self.sphere[i]
        bc,bx,by,bz = sphere.center,sphere.x_axis,sphere.y_axis,sphere.z_axis
        op = ti.math.normalize(p-bc)
        cosb = ti.math.dot(op,bz-bc)
        op_xy = ti.math.normalize(op-(bz-bc)*cosb)
        cosa = ti.math.dot(ti.math.normalize(bx-bc),op_xy)
        v = (np.pi-ti.math.acos(cosb))/np.pi
        u = 0.
        if ti.math.dot(ti.math.normalize(by-bc),op_xy)>0:
            u = ti.math.acos(cosa)/(2.*np.pi)
        else:
            u = 1-ti.math.acos(cosa)/(2.*np.pi)
        return self.get_uv_tex(tex,u,v)

    @ti.func
    def get_object_color(self,ints:Intersection):
        """
        返回纹理颜色,p位置向量,m材质结构,k包围球索引
        l 是三角形索引
        """
        c = ti.Vector([0.,0.,0.])
        m = self.get_material(ints.m)
        if m.tex<0:
            c = m.color
        else:
            if ints.type==PrimitiveType.SPHERE.value:
                c = self.get_sphere_tex(ints.p,ints.i,m.tex)
            else:
                c = self.get_trangle_tex(ints.p,ints.i,m.tex)
        return c

    @ti.func
    def environment_sample(self,ray:Ray):
        """
        取环境光
        """
        c = self.environment_color[None]
        if self.sky_i>=0:
            light = self.light[self.sky_i]
            isc,t,p,n = cross_ray_sphere(ray,self.sphere[light.offset])
            if isc and light.tex>=0:
                sqr_d = SKY_SPHERE_RADIUS_SQR
                c =  self.get_sphere_tex(p,light.offset,light.tex)
                c *= SKY_POWER_SCALE*light.color*light.power/(sqr_d * light.pdf_v)
        return c
    
    @ti.func
    def random_light(self):
        """
        根据cdf选择灯
        """
        r = ti.random()
        sel = -1
        for i in range(self.light_n):
            light = self.light[i]
            if r<light.cdf_v:
                sel = i
                break
        return sel

    @ti.func
    def light_sample(self,ints:Intersection):
        """
        对光进行采样
        返回 v光线是否能到达,le辐射强度和,d方向
        """
        v = 0
        sq = 0.
        p = ti.Vector([0.,0.,0.])
        le = ti.Vector([0.,0.,0.])
        d = ti.Vector([0.,0.,0.])
        sel = self.random_light()
        if sel>=0:
            light = self.light[sel]
            if light.type==LightType.POINT_LIGHT.value:
                p,le,sq = point_light(light,ints)
            elif light.type==LightType.SKY_LIGHT.value:
                p,le,sq = sky_light(light,ints,self)
            elif light.type==LightType.DIRECT_LIGHT.value:
                p,le,sq = direct_light(light,ints)
            elif light.type==LightType.AREA_LIGHT.value:
                p,le,sq = area_light(light,ints,self)
        d = ti.math.normalize(p-ints.p)
        b,intss = self.intersection(Ray(ints.p+ints.n*EPS,d))
        m = self.get_material(intss.m)
        if b and intss.t*intss.t<sq and m.type==MaterialType.BTDF.value:
            v = 1
            le *= 0.618*m.color #这里没有考虑折射对光源的影响
        else:
            v = 1 if not b or intss.t*intss.t>sq else 0
        return v,le,d
